package ratelimit

import (
	"context"
	"go.uber.org/zap"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"jk-time/webook/pkg/limiter"

	"google.golang.org/grpc"
)

type InterceptorBuilder struct {
	limiter limiter.Limiter
	key     string
}

func NewInterceptorBuilder(limiter limiter.Limiter) *InterceptorBuilder {
	return &InterceptorBuilder{limiter: limiter, key: "server_limit"}
}

func (b *InterceptorBuilder) BuildServerInterceptor() grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo,
		handler grpc.UnaryHandler) (resp any, err error) {
		limited, err := b.limiter.Limit(ctx, b.key)
		if err != nil {
			zap.L().Error("判定限流出现问题", zap.Error(err))
			return nil, status.Errorf(codes.ResourceExhausted, "触发限流")
		}
		if limited {
			return nil, status.Errorf(codes.ResourceExhausted, "触发限流")
		}
		return handler(ctx, req)
	}
}
func (b *InterceptorBuilder) BuildServerInterceptorV1() grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo,
		handler grpc.UnaryHandler) (resp any, err error) {
		limited, err := b.limiter.Limit(ctx, b.key)
		if err != nil {
			zap.L().Error("判定限流出现问题", zap.Error(err))
			return nil, status.Errorf(codes.ResourceExhausted, "触发限流")
		}
		if limited {
			ctx = context.WithValue(ctx, "downgrade", "true")
			return handler(ctx, req)
		}
		return handler(ctx, req)
	}
}
func (b *InterceptorBuilder) BuildClientInterceptor() grpc.UnaryClientInterceptor {
	return func(ctx context.Context, method string, req, reply any,
		cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
		limited, err := b.limiter.Limit(ctx, b.key)
		if err != nil {
			zap.L().Error("判定限流出现问题", zap.Error(err))
			return status.Errorf(codes.ResourceExhausted, "触发限流")
		}
		if limited {
			return status.Errorf(codes.ResourceExhausted, "触发限流")
		}
		return invoker(ctx, method, req, reply, cc, opts...)
	}
}
